#include<cstdlib>
#include<cstring>
#include<iostream>
#include<sstream>
#include<fstream>
#include<math.h>
#include<vector>
#include<gsl/gsl_cdf.h>
#include"../common/main.hpp"
#include"../common/analyzer.hpp"
#include"../common/io.hpp"
#include"stepwise.hpp"


using namespace std;
typedef unsigned int uint;

Stepwise::Stepwise(){
}

void Stepwise::init(const ptree & pt){
  settings = new stepwise_settings_t;
  settings->pedfile = pt.get<string>("pedfile");
  settings->snpfile = pt.get<string>("snpfile");
  settings->genofile = pt.get<string>("genofile");
  settings->covariatedatafile = pt.get<string>("covariates.datafile");
  settings->covariateselectionfile = pt.get<string>("covariates.selectionfile");
  string enable = pt.get<string>("db_enable");
  settings->use_db = !enable.compare("true")?true:false;
  settings->db_host = pt.get<string>("db_host");
  settings->db_user = pt.get<string>("db_user");
  settings->db_pw = pt.get<string>("db_pw");
  settings->db_db = pt.get<string>("db_db");
  settings->db_table = pt.get<string>("db_table");
  string t =  pt.get<string>("impute_missing");
  settings->impute_missing = !t.compare("true")?true:false;
  settings->penalty_type = pt.get<string>("penalty_type");
  settings->penalty_file = pt.get<string>("penalty_file");
  io = new IO();
  io->connectToDB(settings->db_host.data(),settings->db_user.data(),settings->db_pw.data(),settings->db_db.data());
  cerr<<"Reading input\n";
  io->readCovariates(settings->covariatedatafile.data(),settings->covariateselectionfile.data(),cmap);
  io->readPlinkFiles(settings->snpfile.data(),settings->pedfile.data(),settings->genofile.data(),data);
  init();
}

Stepwise::~Stepwise(){
  delete[] workingdesignMat;
  delete[] fulldesignMat;
  delete[] fullaff;
  delete[] snparr;
  delete[] covarr;
  delete[] snp_in_model;
  delete[] covarmat;
  delete[] expgeno;
  cerr<<"Stepwise object deleted\n";
}

float Stepwise::penalty(region_t penalties,string rs){
  map<string,float>::iterator it = penalties.find(rs);  
  if (it==penalties.end()){
    throw "Not in penalty map";
  }
  return it->second;
}
 
void Stepwise::initRegion(region_t r){
  totalsnps = r.size();
  int i=0;
  for(region_t::iterator it = r.begin();it!=r.end();it++){
    snparr[i] = it->first; 
    //snp_in_model[i] = true;
    snp_in_model[i] = false;
    ++i;
  }
  int & n = data.totalpersons;
  bool complete_case = true;
  if (complete_case){
    totalobs = 0;
    float * tempgenomat = new float[totalsnps*n]; 
    for (int j = 0;j<totalsnps;++j){
      //cerr<<"SNP index: "<<io->snpindex(snparr[j])<<" for "<<snparr[j]<<endl;
      int index = io->snpindex(snparr[j],data);
      if (index==-1 && settings->use_db){
        //cerr<<"Attempting to fetch imputed SNP: "<<snparr[j]<<endl;
        if (!io->fetchimputedgeno(settings->db_table.data(),snparr[j],tempgenomat+(j*n),n)){
          cerr<<"Failed to locate "<<snparr[j]<<". Aborting.\n";
          throw "SNP not found.";
        }
      }else{
        io->fetchgeno(data.genomatrix, data.veclen, index, tempgenomat+(j*n),data.totalpersons);
      }
      int a = 0,c=0;
      for(int i = 0;i<n;++i){
        if(tempgenomat[j*n+i]!=IO::MISSING_GENO){
          a+=static_cast<int>(tempgenomat[j*n+i]);
          c+=2;
        }
      }
      expgeno[j] = 2.*a/c;
      //cerr<<"Exp geno  of SNP "<<j<<": "<<expgeno[j]<<endl;
    }
    ofstream ofs("demo.txt");
    ofs<<"aff";
    for(int i=0;i<totalcovs;++i){
      ofs<<" c"<<i;
    }
    for(int i=0;i<totalsnps;++i){
      ofs<<" "<<snparr[i];
    }
    ofs<<endl;
    for (int i = 0;i<n;++i){
      int missing = 0; 
      // DISEASE STATUS
      missing+=data.affection[i]==IO::AFF_UNDEF;
//cout<<i<<" "<<missing<<data.affection[i]<<endl;
      fullaff[totalobs] = data.affection[i]==IO::AFF_CASE?1:0;
      ofs<<fullaff[totalobs];
      // INTERCEPT
      workingdesignMat[totalobs * max_rank]  
      = fulldesignMat[totalobs * max_rank] = 1.;
      int offset = 1;
      // COVARIATES
      for(int j = 0;j<totalcovs;++j){
        int ind = totalobs * max_rank + offset + j;
        workingdesignMat[ind] 
        = fulldesignMat[ind] = covarmat[j*n+i];
        if (fulldesignMat[ind]==IO::AFF_UNDEF){
          ofs<<" NA";
        }else{
          ofs<<" "<<fulldesignMat[ind];
        }
        missing+=(fulldesignMat[ind]==IO::AFF_UNDEF);
      }
      offset+=totalcovs;
      // PREDICTORS
      for(int j = 0;j<totalsnps;++j){
        int ind = totalobs * max_rank + offset + j;
        if (settings->impute_missing && tempgenomat[j*n+i]==IO::MISSING_GENO){
          tempgenomat[j*n+i] = expgeno[j];
        }
        workingdesignMat[ind] 
        = fulldesignMat[ind] = tempgenomat[j*n+i];
        if (fulldesignMat[ind]==IO::MISSING_GENO){
          ofs<<" NA";
        }else{
          ofs<<" "<<fulldesignMat[ind];
        }
        missing+=(fulldesignMat[ind]==IO::MISSING_GENO);
      }
      if (missing==0){
        ++totalobs;
      }else{
        //cerr<<"Missing "<<missing<<" values at row "<<i<<endl;
      }
      ofs<<endl;
    }
    ofs.close();
    //exit(0);
    delete[] tempgenomat;
  }
  
  cout<<"Initialized stepwise with base model:\n";
  if (!testmodel(-1,current_logL,true)){
    cerr<<"Exiting as the full model does not converge\n";
    exit(1); 
  }
  //exit(0);
  //snp_in_model[1] = false; 
  //updateDesignMatrix();
}

bool Stepwise::testmodel(int flip_index,double & logL,bool verbose){
  string snplist[totalsnps];
  int k=0;
  for(int j=0;j<totalsnps;++j){
    if (snp_in_model[j]) snplist[k++] = snparr[j];
  }
  
  for(int i=0;i<totalobs;++i){
    currentrank = 1+totalcovs;
    int offset = currentrank;
    for(int j=0;j<totalsnps;++j){
      bool in_model = flip_index==(j)?!snp_in_model[j]:snp_in_model[j];
      if (in_model){
        workingdesignMat[i*max_rank+currentrank] 
        = fulldesignMat[i*max_rank+offset+j];
        //cerr<<" "<<workingdesignMat[i*max_rank+currentrank];
        ++currentrank;
      }
    }
    //cerr<<endl;
  }
  //exit(0);
  double betas[currentrank];
  double var[currentrank*currentrank];
  memset(betas,0,sizeof(betas));
  bool status;
  status = fitModel(logL,fullaff,workingdesignMat,betas,var,totalobs,max_rank,currentrank);
  if (verbose){
    cout<<"variable\tBeta\tSE\tWald_test"<<endl;
    for(int j=0;j<currentrank;++j){
      string rs;
      if (j>=1+totalcovs){
        rs = snplist[j-(1+totalcovs)];
        //rs = snparr[j-1+totalcovs];
      }else{
        if (j==0) rs="intercept";
        else rs = covarr[j-1];
      }
      cout<<rs<<"\t"<<betas[j]<<"\t"<<sqrt(var[j*currentrank+j]);
      double chi = betas[j]*betas[j]/var[j*currentrank+j];
      double pval = (1.-gsl_cdf_chisq_P(chi,1));
      cout<<"\t"<<pval;
      cout<<endl;
    }
    cout<<"Total observations: "<<totalobs<<endl;
  }
  return status;
}

void Stepwise::run(){
  int regionid=1;
  for(list<region_t>::iterator it = region_list.begin();it!=region_list.end();it++){
    cout<<"*** FITTING REGION "<<regionid<<" ***\n";
    region_t r = *it;
    initRegion(r);
   // cerr<<"got here\n";
   // exit(0);
    fitRegion(r);
    ++regionid;
  }
}

void Stepwise::init(){
  max_obs = data.totalpersons;
  ifstream ifs(settings->penalty_file.data());
  if (!ifs.is_open()){
    throw "Cannot open penalty file";
  }
  string line;
  totalsnps = 0;
  string lastregion="";
  region_t penalties;
  max_rank = 0;
  while(getline(ifs,line)){
    string region,rs;
    float penalty;
    istringstream iss(line);
    iss>>region>>rs>>penalty;
    if (region.compare(lastregion)){
      if (lastregion.compare("")){
        region_t newregion;
        newregion = penalties;
        if ((penalties.size()+1)>max_rank) max_rank = penalties.size()+1;
        region_list.push_back(newregion);
        cerr<<"Added region "<<lastregion<<" of size "<<penalties.size()<<endl;
      }
      penalties.clear();
    }
    //cerr<<"pushing "<<rs<<endl;
    penalties[rs] = penalty;
    ++totalsnps;
    lastregion = region;
  }
  region_t newregion;
  newregion = penalties;
  if ((penalties.size()+1)>max_rank) max_rank = penalties.size()+1;
  cerr<<"Added final region "<<lastregion<<" of size "<<penalties.size()<<endl;
  region_list.push_back(newregion);
  ifs.close();

  // populate with covariates
  //covar_map_t cmap = io->getcovars();
  totalcovs = cmap.size();
  // allocate data structures
  max_rank+=totalcovs;

  covarmat = new float[max_obs * totalcovs];
  covarr = new string[totalcovs];
  // load covariates 
  uint covid=0;
  for(covar_map_t::iterator it = cmap.begin();it!=cmap.end();it++){
    vector<float> c = it->second;
    covarr[covid] = it->first;
    for(uint person=0;person<c.size();++person){
      covarmat[covid * max_obs + person] = c[person];
    }
    ++covid;
  }

  snparr = new string[max_rank];
  expgeno = new float[max_rank];
  snp_in_model = new bool[max_rank];
  fulldesignMat = new float[max_obs * max_rank];
  workingdesignMat = new float[max_obs * max_rank];
  fullaff = new int[max_obs];
}

void Stepwise::fitRegion(region_t r){
  bool converged = false;
  int iter=1;
  while(!converged){
    // do nothing case
    cout<<"stepwise_iteration\t"<<iter++<<endl;
    cout<<"intercept";
    for(int i=0;i<totalsnps;++i){
      if (snp_in_model[i]) cout<<" + "<<snparr[i];
    }
    cout<<endl;
    ModelChange mc0;
    mc0.rsindex = -1;
    mc0.action = 0;
    mc0.logL = current_logL;
    mc0.score = 2.*current_logL;
    //cout<<"Candidate\tAction\tlogL\tPenalized_score\n";
    //cout<<"none\t"<<mc0.action<<"\t"<<mc0.logL<<"\t"<<mc0.score<<endl;
    ModelChange best_mc = mc0;

    for(int i=0;i<totalsnps;++i){
      ModelChange mc;
      mc.rsindex=i;
      if (!snp_in_model[i]){
        mc.action = 1;
      }else{
        mc.action = -1;
      }
      if(testmodel(i,mc.logL,false)){
        mc.score=2.*mc.logL - 1. * mc.action * penalty(r,snparr[mc.rsindex]);
        //cout<<snparr[mc.rsindex]<<"\t"<<mc.action<<"\t"<<mc.logL<<"\t"<<mc.score<<endl;
        if (mc.score>best_mc.score){
          best_mc = mc;
          //cout<<"so far Best model change:\n"<<best_mc.rs<<"\t"<<best_mc.action<<"\t"<<best_mc.logL<<endl;
        }
      }
    }
    string bestsnp = best_mc.rsindex>=0?snparr[best_mc.rsindex]:"no_snp";
    cout<<"Best model change: "<<bestsnp<<"\t"<<best_mc.action<<"\t"<<best_mc.logL<<"\t"<<best_mc.score<<endl;
    if (best_mc.score>mc0.score){
      snp_in_model[best_mc.rsindex] = !snp_in_model[best_mc.rsindex];
      current_logL = best_mc.logL;
    }else{
      converged = true;
    }
  }
  cout<<"Final model:\n";
  testmodel(-1,current_logL,true);
}

